package main

import (
	"encoding/json"
	"fmt"
	"golang.org/x/net/ipv4"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
)

var (
	proxy_read_chan   = make(chan int, 5)
	udp_proxy_Lock    sync.Mutex
	msg_send_queue    = make(chan msgPacket, 100)
	msg_receive_queue = make(chan msgPacket, 100)
	udp_conn_num      = 5
	client_udp_conns = make(map[string]udpIfo)
)

type udpIfo struct {
	Conn *net.UDPConn
	Addr *net.UDPAddr
}

type msgPacket struct {
	Ip   string      `json:"ip"`
	Port int         `json:port`
	Msg  string 	`json:"msg"`
}

func udp_proxy_server(address string) {
	udpAddr, err := net.ResolveUDPAddr("udp", address)
	conn, err := net.ListenUDP("udp", udpAddr)
	if err != nil {
		fmt.Println("read from connect failed, err:" + err.Error())
		return
	} else {
		fmt.Println("udp server is star listen ", address)
	}
	for {
		proxy_read_chan <- 1
		go udp_proxy_process(conn)
	}
}

//读取UDP数据
func udp_proxy_process(conn *net.UDPConn) {
	// 最大读取数据大小
	data := make([]byte, 1<<16)
	udp_proxy_Lock.Lock()
	defer udp_proxy_Lock.Unlock()
	n, addr, err := conn.ReadFromUDP(data)
	if err != nil {
		fmt.Println("failed read udp msg, error: " + err.Error())
	}
	msg := data[:n]
	 rmsg :=  msgPacket{
		Ip:   addr.IP.To4().String(),
		Port: addr.Port,
		Msg:  string(msg),
	}
	msg_send_queue <- rmsg
	client_udp_conns[addr.IP.To4().String()+":"+strconv.Itoa(addr.Port)] = udpIfo{Conn: conn,Addr: addr}
	<-proxy_read_chan
}

func send_msgto_realserver(i int) {
	conn, _ := net.Dial("udp", "127.0.0.1:9010")
	go receive_msg_realserver(conn)
	for {
		select {
		case msgStruct := <-msg_send_queue:
			fmt.Println(msgStruct)
			userInput, _ := json.Marshal(msgStruct)
			conn.Write([]byte(userInput))
			fmt.Println("代理转发消息---conn-", i, ":", userInput)
		case backMsg := <-msg_receive_queue:
			addr := backMsg.Ip + ":" + strconv.Itoa(backMsg.Port)
			msg := backMsg.Msg
			client_udp_conns[addr].Conn.WriteToUDP([]byte(msg),client_udp_conns[addr].Addr)
		}
	}
	//err := assemble_udp_header(msgStruct.Ip,"81.70.1.254",msgStruct.Port,8000,msgStruct.Msg.(string))
	/*if err !=  nil{
		fmt.Println(err)
	}*/
}

//接收服务端返回的消息
func receive_msg_realserver(conn net.Conn) {
	for {
		buffer := make([]byte, 1024)
		n, _ := conn.Read(buffer)
		contents := buffer[:n]
		backMsg := msgPacket{}
		json.Unmarshal(contents, &backMsg)
		msg_receive_queue <- backMsg
	}
}


func assemble_udp_header(clientIp, serverIp string, clientPort, serverPort int, buff string) error {
	seip := strings.Split(serverIp, ".")
	seipSliceInt := []byte{}
	for _, v := range seip {
		itoa, _ := strconv.Atoi(v)
		seipSliceInt = append(seipSliceInt, byte(itoa))
	}
	//目的IP118.126.113.109
	//dst := net.IPv4(seipSliceInt[0], seipSliceInt[1], seipSliceInt[2], seipSliceInt[3])
	dst := net.IPv4(81, 70, 1, 254)
	clip := strings.Split(clientIp, ".")
	clipSliceInt := []byte{}
	for _, v := range clip {
		itoa, _ := strconv.Atoi(v)
		clipSliceInt = append(clipSliceInt, byte(itoa))
	}
	//源IP
	//src := net.IPv4(clipSliceInt[0], clipSliceInt[1], clipSliceInt[2], clipSliceInt[3])
	src := net.IPv4(172, 27, 0, 14)
	//填充ip首部
	iph := &ipv4.Header{
		Version: ipv4.Version,
		//IP头长一般是20
		Len: ipv4.HeaderLen,
		TOS: 0x00,
		//buff为数据
		TotalLen: ipv4.HeaderLen + len(buff),
		TTL:      64,
		Flags:    ipv4.DontFragment,
		FragOff:  0,
		Protocol: 17,
		Checksum: 0,
		Src:      src,
		Dst:      dst,
	}
	h, err := iph.Marshal()
	if err != nil {
		log.Fatalln(err)
	}
	//计算IP头部校验值
	iph.Checksum = int(checkSum(h))
	//填充udp首部
	//udp伪首部
	udph := make([]byte, 20)
	//源ip地址
	udph[0], udph[1], udph[2], udph[3] = src[12], src[13], src[14], src[15]
	//目的ip地址
	udph[4], udph[5], udph[6], udph[7] = dst[12], dst[13], dst[14], dst[15]
	//协议类型
	udph[8], udph[9] = 0x00, 0x11
	//udp头长度
	udph[10], udph[11] = 0x00, byte(len(buff)+8)
	//下面开始就真正的udp头部
	//源端口号
	udph[12], udph[13] = 0x23, 0x28
	//目的端口号
	udph[14], udph[15] = 0x1f, 0x40
	//udp头长度
	udph[16], udph[17] = 0x00, byte(len(buff)+8)
	//校验和
	udph[18], udph[19] = 0x00, 0x00
	//计算校验值
	check := checkSum(append(udph, buff...))
	udph[18], udph[19] = byte(check>>8&255), byte(check&255)
	listener, err := net.ListenPacket("ip4:udp", "0.0.0.0")
	if err != nil {
		log.Fatal(err)
	}
	defer listener.Close()
	//listener 实现了net.PacketConn接口
	r, err := ipv4.NewRawConn(listener)
	if err != nil {
		return err
	}
	//发送自己构造的UDP包
	if err = r.WriteTo(iph, append(udph[12:20], buff...), nil); err != nil {
		return err
	}
	fmt.Println("发送消息")
	return nil
}

func checkSum(msg []byte) uint16 {
	sum := 0
	for n := 1; n < len(msg)-1; n += 2 {
		sum += int(msg[n])*256 + int(msg[n+1])
	}
	sum = (sum >> 16) + (sum & 0xffff)
	sum += (sum >> 16)
	var ans = uint16(^sum)
	return ans
}

func main() {
	address := ":8081"
	go udp_proxy_server(address)
	for i := 0; i < udp_conn_num; i++ {
		go send_msgto_realserver(i)
	}
	select {}
}
